from .strategies import *
from .models import VideoConfig, FrameData
from .processor import VideoProcessor
from .saver import KeyframeSaver

from skimage.metrics import structural_similarity as ssim
from typing import List, Set
import os
import cv2
import numpy as np
from src.config_reader import ConfigReader

class FrameExtractor:
    def __init__(self):
        config = ConfigReader()
        base_dir = config.get('output.base_dir_path')
        self.output_base_dir = os.path.join(base_dir, 
                                            config.get('output.image_output_base_path'))
        self.video_src_path = os.path.join(base_dir, 
                                           config.get('output.video_src_path'))
        
        self.config = VideoConfig(
            len_window=config.get('processing.len_window'),
            use_thresh=config.get('processing.use_thresh'),
            thresh=config.get('processing.thresh'),
            use_top_order=config.get('processing.use_top_order'),
            num_top_frames=config.get('processing.num_top_frames'),
            use_local_maxima=config.get('processing.use_local_maxima'),
            ssim_threshold=config.get('processing.ssim_threshold'),
            start_time=config.get('processing.start_time'),
            end_time=config.get('processing.end_time')
        )

    def process(self, video_name:str):
        # 处理视频获取帧数据并检测关键帧，最后保存关键帧
        video_path = os.path.join(self.video_src_path, video_name)
        print(f"视频路径: {video_path}")
        with VideoProcessor(video_path, self.config.start_time, self.config.end_time) as processor:
            frames = self._process_frames(processor)
        
        # 检测关键帧
        strategies = self._get_strategies() # 根据配置选择策略
        keyframes = set()
        for strategy in strategies:
            keyframes.update(strategy.detect(frames))
        
        # 保存关键帧
        saver = KeyframeSaver(self._get_output_dir(video_path), self.config.ssim_threshold)
        saver.save(video_path, keyframes)

    def _process_frames(self, processor: VideoProcessor) -> List[FrameData]:
        frames = []
        prev_frame = None
        for i, curr_frame in enumerate(processor.get_frames()):
            # 仅在计算差异时转换为LUV
            curr_luv = cv2.cvtColor(curr_frame, cv2.COLOR_BGR2LUV)
            if prev_frame is not None:
                prev_luv = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2LUV)
                diff = cv2.absdiff(curr_luv, prev_luv)
                diff_mean = np.sum(diff) / (diff.shape[0] * diff.shape[1])
                frames.append(FrameData(i, diff_mean))
            prev_frame = curr_frame  # 保存原始BGR帧供下一轮比较
        return frames

    def _get_strategies(self) -> List[KeyframeStrategy]:
        # 根据配置选择并初始化关键帧检测策略
        strategies = []
        if self.config.use_thresh:
            strategies.append(ThresholdStrategy(self.config.thresh))
        if self.config.use_top_order:
            strategies.append(TopFramesStrategy(self.config.num_top_frames))
        if self.config.use_local_maxima:
            strategies.append(LocalMaximaStrategy(self.config.len_window))
        return strategies

    def _get_output_dir(self, video_path: str) -> str:
        # 根据视频路径生成输出目录
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        return os.path.join(self.output_base_dir, video_name)